package com.xunjing.order.config.sharding;

import com.xunjing.common.core.utils.SpringUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.ShardingAlgorithm;
import org.apache.shardingsphere.core.rule.TableRule;
import org.apache.shardingsphere.core.strategy.route.ShardingStrategy;
import org.apache.shardingsphere.core.strategy.route.complex.ComplexShardingStrategy;
import org.apache.shardingsphere.core.strategy.route.hint.HintShardingStrategy;
import org.apache.shardingsphere.core.strategy.route.standard.StandardShardingStrategy;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.context.ShardingRuntimeContext;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.springframework.core.env.Environment;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import javax.annotation.Resource;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.sql.*;
import java.time.YearMonth;
import java.time.format.DateTimeFormatter;
import java.util.*;

/**
 * @Author ZYT
 * @Date 2024-03-29 18:11
 * @Version 1.0
 * @Describe 分表工具类
 */
@Slf4j
@Component
public class ShardingAlgorithmReload {

    @Resource
    private ShardingDataSource shardingDataSource;

    private ShardingRuntimeContext runtimeContext;

    /**
     * 表分片符号，例：t_user_202201 中，分片符号为 "_"
     */
    private static final String TABLE_SPLIT_SYMBOL = "_";

    /**
     * 重载表缓存
     */
    public void tableNameCacheReloadAll() {
        ShardingRuntimeContext runtimeContext = getRuntimeContext();

        List<TableRule> tableRuleList = (List<TableRule>) runtimeContext.getRule().getTableRules();
        for (TableRule tableRule : tableRuleList) {
            String nodeName = tableRule.getActualDatasourceNames().stream().findFirst().get();
            Set<String> tablesInDBSet = queryTables(tableRule.getLogicTable());
            refreshTableRule(tableRule, nodeName, tablesInDBSet);
            //缓存全节点（指定节点）、可以不用
            //refreshShardingAlgorithm(tableRule, null);
        }
    }

    //刷新表规则
    protected void refreshTableRule(TableRule tableRule, String nodeName, Set<String> tablesInDBSet) {
        // sharding缓存的表名
        Set<String> actualTableSets = getActualTables(tableRule);
        // 刷新
        if (!tableContrast(actualTableSets, tablesInDBSet)) {
            List<String> tableList = new ArrayList<>(tablesInDBSet);
            setDatasourceToTablesMap(tableRule, nodeName, tableList);
        }

    }

    private boolean tableContrast(Set<String> actualTableSets, Set<String> tablesInDBSet) {
        if (actualTableSets == null || tablesInDBSet == null) {
            return false;
        }
        if (actualTableSets.size() != tablesInDBSet.size()) {
            return false;
        }
        return actualTableSets.containsAll(tablesInDBSet);
    }

    protected void refreshShardingAlgorithm(TableRule tableRule, String nodeName) {
        // 获取分库分表时真正使用的表名
        Map<String, Set<String>> datasourceToTablesMap = getDatasourceToTablesMap(tableRule);
        Set<String> tables = datasourceToTablesMap.get(nodeName);
        ShardingStrategy shardingStrategy = tableRule.getTableShardingStrategy();
        if (shardingStrategy instanceof ComplexShardingStrategy) {
            ShardingAlgorithm algorithm = getObjectField(shardingStrategy, "shardingAlgorithm");
            setValueToBaseAlgorithm(tableRule, algorithm, nodeName, tables);
        } else if (shardingStrategy instanceof HintShardingStrategy) {
            ShardingAlgorithm algorithm = getObjectField(shardingStrategy, "shardingAlgorithm");
            setValueToBaseAlgorithm(tableRule, algorithm, nodeName, tables);
        } else if (shardingStrategy instanceof StandardShardingStrategy) {
            ShardingAlgorithm preciseAlgorithm = getObjectField(shardingStrategy, "preciseShardingAlgorithm");
            setValueToBaseAlgorithm(tableRule, preciseAlgorithm, nodeName, tables);
            ShardingAlgorithm rangeAlgorithm = getObjectField(shardingStrategy, "rangeShardingAlgorithm");
            setValueToBaseAlgorithm(tableRule, rangeAlgorithm, nodeName, tables);
        }
    }

    private void setValueToBaseAlgorithm(TableRule tableRule, ShardingAlgorithm algorithm, String nodeName, Set<String> tables) {

        if (algorithm != null && algorithm instanceof BaseShardingAlgorithm) {
            BaseShardingAlgorithm baseShardingAlgorithm = (BaseShardingAlgorithm) algorithm;
            baseShardingAlgorithm.setLogicTable(tableRule.getLogicTable());
            baseShardingAlgorithm.setTables(tables);
            baseShardingAlgorithm.setTableRule(tableRule);
            baseShardingAlgorithm.setNodeName(nodeName);
        }
    }

    protected ShardingRuntimeContext getRuntimeContext() {
        try {
            if (runtimeContext == null) {
                Method getRuntimeContextMethod = shardingDataSource.getClass().getDeclaredMethod("getRuntimeContext");
                getRuntimeContextMethod.setAccessible(true);
                runtimeContext = (ShardingRuntimeContext) getRuntimeContextMethod.invoke(shardingDataSource, null);
            }
        } catch (IllegalAccessException e) {
            log.error("因为sharding版本问题", e);
        } catch (InvocationTargetException e) {
            log.error("因为sharding版本问题", e);
        } catch (NoSuchMethodException e) {
            log.error("因为sharding版本问题", e);
        }
        return runtimeContext;
    }

    protected Set<String> getActualTables(TableRule tableRule) {
        Set<String> tables = getObjectField(tableRule, "actualTables");
        return tables == null ? new LinkedHashSet<>() : tables;
    }

    protected void setDatasourceToTablesMap(TableRule tableRule, String nodeName, List<String> newTableList) {
        synchronized (tableRule) {
            // 获取分库分表时真正使用的表名
            Map<String, Set<String>> datasourceToTablesMap = getDatasourceToTablesMap(tableRule);
            Set<String> tables = datasourceToTablesMap.get(nodeName);
            Collections.sort(newTableList);
            tables.clear();
            tables.addAll(newTableList);
        }
    }

    protected Map<String, Set<String>> getDatasourceToTablesMap(TableRule tableRule) {
        Map<String, Set<String>> tablesMap = getObjectField(tableRule, "datasourceToTablesMap");
        return tablesMap == null ? new HashMap<>(0) : tablesMap;
    }

    protected static <T> T getObjectField(Object object, String fieldName) {
        try {
            Field field = object.getClass().getDeclaredField(fieldName);
            field.setAccessible(true);
            return (T) field.get(object);
        } catch (NoSuchFieldException e) {
            log.error("因为sharding版本问题", e);
        } catch (IllegalAccessException e) {
            log.error("因为sharding版本问题", e);
        }
        return null;
    }

    protected Set<String> queryTables(String tableName) {
        Connection conn = null;
        Statement statement = null;
        ResultSet rs = null;
        Set<String> tables = null;
        try {
            conn = shardingDataSource.getConnection();
            statement = conn.createStatement();
            rs = statement.executeQuery("select table_name from information_schema.tables where table_schema ='数据库名' and table_name LIKE '" + tableName + "_%'");
            tables = new LinkedHashSet<>();
            while (rs.next()) {
                tables.add(rs.getString(1));
            }
        } catch (SQLException e) {
            log.error("获取数据库连接失败！", e);
        } finally {
            try {
                if (rs != null) {
                    rs.close();
                }
                if (statement != null) {
                    statement.close();
                }
                if (conn != null) {
                    conn.close();
                }
            } catch (SQLException e) {
                log.error("关闭数据连接失败", e);
            }
        }
        return tables;
    }

    //创建分表
    public boolean createShardingTable(String logicTableName, String preciseTable) {
        // 根据日期判断，当前月份之后分表不提前创建
        String month = preciseTable.replace(logicTableName + TABLE_SPLIT_SYMBOL, "");
        YearMonth shardingMonth = YearMonth.parse(month, DateTimeFormatter.ofPattern("yyyyMM"));
        if (shardingMonth.isAfter(YearMonth.now())) {
            return false;
        }
        //创建锁
        synchronized (logicTableName.intern()) {
            // 缓存中无此表，则建表并添加缓存
            executeSql(Collections.singletonList("CREATE TABLE IF NOT EXISTS `" + preciseTable + "` LIKE `" + logicTableName + "`;"));
        }
        return true;

    }

    //创建表sql执行
    @Transactional
    public void executeSql(List<String> sqlList) {
        Environment env = SpringUtil.getApplicationContext().getEnvironment();
        try (Connection conn = DriverManager.getConnection(env.getProperty("spring.shardingsphere.datasource.master.url"), env.getProperty("spring.shardingsphere.datasource.master.username"), env.getProperty("spring.shardingsphere.datasource.master.password"))) {
            try (Statement st = conn.createStatement()) {
                conn.setAutoCommit(false);
                for (String sql : sqlList) {
                    st.execute(sql);
                }
            } catch (Exception e) {
                conn.rollback();
                log.error(">>>>>>>>>> 【ERROR】数据表创建执行失败，请稍后重试，原因：{}", e.getMessage(), e);
                throw new IllegalArgumentException("数据表创建执行失败，请稍后重试");
            }
        } catch (SQLException e) {
            log.error(">>>>>>>>>> 【ERROR】数据库连接失败，请稍后重试，原因：{}", e.getMessage(), e);
            throw new IllegalArgumentException("数据库连接失败，请稍后重试");
        }
    }
}


